"""文生图应用

该模块提供了一个文生图应用，可以根据用户输入的创意描述，通过大模型优化提示词，
并使用硅基流动平台的可图大模型生成图像。支持提示词重新生成和图片保存功能。
"""

import os
import requests
from datetime import datetime
from openai import OpenAI
from dotenv import load_dotenv
import random

# pip install requests openai python-dotenv


def get_api_key():
    """获取API密钥
    从环境变量中获取硅基流动平台API密钥
    Returns:
        API密钥
    """
    if os.path.exists(".env"):
        load_dotenv(".env")
        api_key = os.environ.get("SILICONFLOW_API_KEY", "")
        return api_key
    else:
        print("未找到.env文件，请确保环境变量已正确配置")
        return None


def get_base_url():
    """获取API基础URL
    Returns:
        API基础URL
    """
    return "https://api.siliconflow.cn/v1"


def ai_chat(prompt: str, **kwargs):
    """调用AI接口进行对话
    Args:
        prompt: 提示词
        **kwargs: 其他参数，包括model、temperature、max_tokens、stream等
    Returns:
        AI响应结果
    """
    api_key = get_api_key()
    if not api_key:
        print("未提供API密钥，请通过参数传入或设置SILICONFLOW_API_KEY环境变量")
        return None

    base_url = get_base_url()
    client = OpenAI(api_key=api_key, base_url=base_url)

    stream = kwargs.get("stream", True)
    response = client.chat.completions.create(
        model=kwargs.get("model", "deepseek-ai/DeepSeek-V3"),
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": f"{prompt}"},
        ],
        temperature=kwargs.get("temperature", 1.0),
        max_tokens=kwargs.get("max_tokens", 4096),
        stream=stream,
    )

    if stream:
        contents = ""
        print("\nAI回复：")
        for chunk in response:
            content = chunk.choices[0].delta.content
            print(content, end="", flush=True)
            contents += content
    else:
        contents = response.choices[0].message.content
        print("\nAI回复：")
        print(contents)

    return contents


def generate_image_prompt(idea: str) -> str:
    """使用AI生成优化的图像提示词
    Args:
        idea: 用户输入的创意描述
    Returns:
        优化后的图像提示词
    """
    prompt = f"""
    请根据以下创意描述，生成一个详细的中文的图像提示词，用于文生图AI模型。
    要求：
    1. 提示词要具体详细，包含场景、物体、氛围、光线等细节
    2. 使用中文描述
    3. 只返回提示词本身，不要包含其他解释性文字
    4. 提示词长度适中，不要过长

    创意描述：{idea}
    """

    try:
        # 调用AI服务生成提示词
        response = ai_chat(prompt=prompt)
        return response.strip()
    except Exception as e:
        raise Exception(f"生成图像提示词失败: {str(e)}")


def generate_image(prompt: str) -> str:
    """调用可图大模型生成图像
    参考文档：https://docs.siliconflow.cn/cn/api-reference/images/images-generations
    Args:
        prompt: 图像提示词
    Returns:
        生成的图像文件路径
    """
    api_key = get_api_key()
    if not api_key:
        raise Exception("未提供API密钥")

    url = f"{get_base_url()}/images/generations"

    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}

    payload = {
        "model": "Kwai-Kolors/Kolors",
        "prompt": prompt,
        "negative_prompt": "",
        "image_size": "1024x1024",
        "batch_size": 1,
        "seed": random.randint(0, 9999999999),
        "num_inference_steps": 20,
        "guidance_scale": 7.5,
    }

    try:
        print("正在生成图像...")
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()

        # 解析响应
        result = response.json()
        if "images" not in result:
            raise Exception(f"API返回了报错状态码: {response.status_code}")

        # 获取图像地址
        image_url = result["images"][0]["url"]

        # 生成文件名
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        image_path = f"generated_image_{timestamp}.png"

        # 保存图像
        with open(image_path, "wb") as f:
            image_data = requests.get(image_url).content
            f.write(image_data)

        print(f"\n图像已保存至: {image_path}")
        return str(image_path)

    except requests.exceptions.RequestException as e:
        raise Exception(f"调用图像生成API失败: {str(e)}")
    except Exception as e:
        raise Exception(f"生成图像失败: {str(e)}")


def main():
    """主函数，处理用户输入并生成图像"""
    print("=== 文生图应用 ===\n")
    print("说明：")
    print("1. 输入创意描述，AI会帮你优化提示词并生成图像")
    print("2. 如果对生成的图像不满意，可以重新生成提示词")
    print("3. 输入'exit'退出程序\n")

    while True:
        try:
            # 获取用户输入的创意描述
            idea = input("\n请输入您的创意描述（输入'exit'退出）：")
            if idea.lower() == "exit":
                print("\n程序已退出")
                break

            while True:
                # 生成优化的提示词
                print("\n正在生成优化的提示词...")
                prompt = generate_image_prompt(idea)

                # 询问用户是否满意
                confirm = input("\n对提示词满意吗？(y/n): ")
                if confirm.lower() == "y":
                    # 生成图像并保存
                    generate_image(prompt)
                    break
                else:
                    print("\n好的，让我重新生成提示词...")

        except Exception as e:
            print(f"\n发生错误: {str(e)}")
            continue


if __name__ == "__main__":
    main()
